# -*- coding: utf-8 -*-
"""
Created on Thu Apr 21 22:47:46 2022

@author: 13503
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data import DataLoader, Dataset
import os
import h5py
import random

import matplotlib.pyplot as plt
import networkx as nx


import dgl
import dgl.function as dglfn
from dgl.data import DGLDataset

class E2VTrainer():
    def __init__(self, graph:dgl.DGLGraph):
        self.graph = graph

    def train(self, epochs:int):
        for _ in range(epochs):
            self.graph.apply_edges(func=node_diff_to_edge("feat", "feat", "sign"))
            self.graph.update_all(message_func=dglfn.u_add_e("feat", "displacement","m"), 
                                  reduce_func=dglfn.sum("m", "feat"))

class SignNetGraph(DGLDataset):
    def __init__(self, 
                 graph_name:str, 
                 dataset_dir:str, 
                 dataset_names:list[str], 
                 feat_dir:str=None,
                 feat_names:list[str]=None)->None:

        if feat_dir is not None:
            if len(feat_names) != len(dataset_names):
                raise IndexError("The length of datesets does not match the length of feature!")
            
            self.feat_paths = []
            for fname in feat_names:
                self.feat_paths.append(os.path.join(feat_dir, fname))
        else:
            self.feat_paths = None

        self.dataset_paths = []
        for dname in dataset_names:
            self.dataset_paths.append(os.path.join(dataset_dir, dname))
        
        self.graphs = []
        super().__init__(name=graph_name, raw_dir=dataset_dir, save_dir=dataset_dir)

    def process(self) -> None:
        # read data from file to generate Dataset
        for idx in range(len(self.dataset_paths)):
            dataset_path = self.dataset_paths[idx]
            data_in_tensor = torch.load(dataset_path).type(torch.float)
            src, dst = torch.nonzero(data_in_tensor, as_tuple=True)
            graph = dgl.graph((src,dst), num_nodes=data_in_tensor.shape[0])
            #graph.edata["sign"] = data_in_tensor[src,dst].reshape((graph.num_edges(),1))
            
            if self.feat_paths is not None:
                feat_path = self.feat_paths[idx]
                graph.ndata["feat"] = torch.load(feat_path).cpu()

            pos_neg = torch.zeros((graph.num_nodes(),2)) # store how many positive and negative links starts from each node
            signs = torch.zeros((graph.num_edges(),1)) # store the sign of each link
            for i in range(len(src)):
                if data_in_tensor[src[i]][dst[i]] == 1:
                    pos_neg[src[i]][0] += 1
                else:
                    pos_neg[src[i]][1] += 1
            print("There are {0} edges in total ".format(graph.num_edges()))
            print("count the number of negative links = {0}".format(torch.sum(pos_neg.index_select(1,torch.tensor([1])))))
            print("count the number of positve links = {0}".format(torch.sum(pos_neg.index_select(1,torch.tensor([0])))))
            for i in range(len(src)):
                if data_in_tensor[src[i]][dst[i]] == 1:
                    signs[i] = 1.0 / pos_neg[src[i]][0] 
                    
                else:
                    signs[i] = -100.0 / pos_neg[src[i]][1]

            graph.ndata["pos_neg"] = pos_neg
            graph.edata["sign"] = signs
            self.graphs.append(graph)


    def __getitem__(self, idx)->None:
        # get one of the grahs and store it onto device
        return self.graphs[idx]
    
    def __len__(self)-> None:
        return len(self.graphs)

    def save(self)->None:
        # save graphs onto disk
        dgl.save_graphs(self.graphs_path, self.graphs)

    def load(self)->None:
        # load processed data from directory `self.save_path`
        self.graphs, _ = dgl.load_graphs(self.graphs_path)

    def has_cache(self)->bool:
        # check whether there are processed data in `self.save_path`
        self.graphs_path = os.path.join(self.save_dir, self.name.split(".")[0] + '_dgl_graph.bin')
        return os.path.exists(self.graphs_path)


class node_diff_to_edge():
    
    def __init__(self, src_kw:str, dst_kw:str, edge_kw:str, decay_rate:float=0.9):
        """
        Collect information from edge-connected nodes to calculate the displacement direction
        """
        self.src_kw = src_kw
        self.dst_kw = dst_kw
        self.edge_kw = edge_kw
        self.decay_rate = decay_rate

    def __call__(self, edge):            
        edge.data[self.edge_kw] = edge.data[self.edge_kw] * self.decay_rate
        return {"displacement": edge.data[self.edge_kw]  * (edge.src[self.src_kw] - edge.dst[self.dst_kw])}

if __name__ == "__main__":
    adj = torch.tensor([[0, 1, 0, 0, 0, 0],
                        [0, 0, 1, 0, 0, 1],
                        [0, 0, 0, 0, -1, 0],
                        [0, 0, 1, 0, -1, 0],
                        [0, 0, 0, 0, 0, 1],
                        [1, 0, 0, 0, 0, 0]])
    torch.save(adj, os.path.join("dataset", "adj.pt"))

    feats = torch.tensor([
        [0.10, 0.20],
        [0.20, 0.90],
        [0.85, 0.60],
        [0.35, 0.62],
        [0.73, 0.22],
        [0.45, 0.08]
    ])
    torch.save(feats, os.path.join("rep_cache", "feats.pt"))

    graph_train_dataset = SignNetGraph("show_network", 
                              "dataset", 
                              ["adj.pt"],
                              "rep_cache", 
                              ["feats.pt"])

    point_label = ["0", "1", "2", "3", "4", "5"]
    e2v = E2VTrainer(graph_train_dataset[0])
    
    plt.figure()
    x = list(graph_train_dataset[0].ndata["feat"].index_select(1,torch.tensor([0])))
    y = list(graph_train_dataset[0].ndata["feat"].index_select(1,torch.tensor([1])))
    plt.scatter(x, y, marker="o")
    for i in range(len(point_label)):
        plt.annotate(text=point_label[i], xy=(x[i],y[i]))
    plt.show()
    
    for k in range(10):
        e2v.train(1)
        plt.figure()
        x = list(graph_train_dataset[0].ndata["feat"].index_select(1,torch.tensor([0])))
        y = list(graph_train_dataset[0].ndata["feat"].index_select(1,torch.tensor([1])))
        plt.scatter(x, y, marker="o")
        for i in range(len(point_label)):
            plt.annotate(text=point_label[i], xy=(x[i],y[i]))
        plt.show()